use spin::Mutex;
use alloc::collections::BTreeMap;
use alloc::sync::Arc;
use alloc::vec::Vec;
use bitflags::bitflags;
use lazy_static::*;
use riscv::register::satp;

use crate::config::{PAGE_SIZE, TRAMPOLINE, MEMORY_END, USER_STACK_SIZE, TRAP_CONTEXT};
use crate::link_symbols::*;
use super::address::*;
use super::frame_allocator::*;
use super::page_table::*;

// 地址空间是一系列有关联的逻辑段
// 这种关联一般是指这些逻辑段属于一个运行的程序
// 用来表明正在运行的应用所在执行环境中的可访问内存空间
// 在这个内存空间中，包含了一系列的不一定连续的逻辑段。
pub struct MemorySet {
    page_table: PageTable, // 该地址空间的多级页表 page_table
    areas: Vec<MapArea>,   // 逻辑段 MapArea 组
    // PageTable 下 挂着所有多级页表的节点所在的物理页帧
    // 而每个 MapArea 下则挂着对应逻辑段中的数据所在的物理页帧
    // 当一个地址空间 MemorySet 生命周期结束后， 这些物理页帧都会被回收
}

impl MemorySet {
    // 新建一个空的地址空间
    pub fn new_bare() -> Self {
        Self {
            page_table: PageTable::new(),
            areas: Vec::new(),
        }
    }

    // 在当前地址空间中插入一个新的逻辑段 map_area
    // 如果映射方式为 Framed(在 copy_data 中通过断言确保),
    // 还可以可选地在被映射到的物理页帧上写下一些初始化数据
    fn push(&mut self, mut map_area: MapArea, data: Option<&[u8]>) {
        map_area.map(&mut self.page_table); // 维护地址空间的多级页表 page_table 记录的虚拟页号到页表项的映射关系
        if let Some(data) = data { // 写入数据
            map_area.copy_data(&mut self.page_table, data);
        }
        self.areas.push(map_area);  // 将该逻辑段挂载到当前地址空间中, 对其生命周期进行管理
    }

    // 在当前地址空间插入一个 Framed 方式映射到 物理内存的逻辑段
    // 该方法的调用者要保证同一地址空间内的任意两个逻辑段不能存在交集
    // 该方法是对私有接口 push 的封装
    pub fn insert_framed_area(
        &mut self,
        start_va: VirtAddr, end_va: VirtAddr, permission: MapPermission
    ) {
        self.push(MapArea::new(
            start_va,
            end_va,
            MapType::Framed,
            permission,
        ), None);
    }

    pub fn find_vpn(&self, vpn: VirtPageNum) -> bool {
        self.page_table.find_vpn(vpn)
    }

    pub fn munmap_area(&mut self, start_va: VirtAddr, end_va: VirtAddr) -> bool {
        let start_vpn: VirtPageNum = start_va.floor();
        let end_vpn: VirtPageNum = end_va.ceil();
        let range = VPNRange::new(start_vpn, end_vpn);

        for vpn in range {
            if self.page_table.find_vpn(vpn) {
                self.page_table.unmap(vpn);
            } else {
                return false;
            }
        }
        return true;
    }

    pub fn translate(&self, vpn: VirtPageNum) -> Option<PageTableEntry> {
        self.page_table.translate(vpn)
    }

    pub fn token(&self) -> usize {
        self.page_table.token()
    }

    // 为跳板在多级页表中建立映射
    fn map_trampoline(&mut self) {
        self.page_table.map(
            VirtAddr::from(TRAMPOLINE).into(),
            PhysAddr::from(strampoline as usize).into(),
            PTEFlags::R | PTEFlags::X,
        );
    }

    // 用于生成内核的地址空间
    pub fn new_kernel() -> Self {
        let mut memory_set = Self::new_bare();
        // 为跳板在多级页表中建立映射
        memory_set.map_trampoline();
        // 为内核各个段建立映射
        info!(".text [{:#x}, {:#x})", stext as usize, etext as usize);
        info!(".rodata [{:#x}, {:#x})", srodata as usize, erodata as usize);
        info!(".data [{:#x}, {:#x})", sdata as usize, edata as usize);
        info!(".bss [{:#x}, {:#x})", sbss_with_stack as usize, ebss as usize);
        
        info!("mapping .text section");
        memory_set.push(MapArea::new(
            (stext as usize).into(),
            (etext as usize).into(),
            MapType::Identical,
            MapPermission::R | MapPermission::X,
        ), None);
        
        info!("mapping .rodata section");
        memory_set.push(MapArea::new(
            (srodata as usize).into(),
            (erodata as usize).into(),
            MapType::Identical,
            MapPermission::R,
        ), None);

        info!("mapping .data section");
        memory_set.push(MapArea::new(
            (sdata as usize).into(),
            (edata as usize).into(),
            MapType::Identical,
            MapPermission::R | MapPermission::W,
        ), None);

        info!("mapping .bss section");
        memory_set.push(MapArea::new(
            (sbss_with_stack as usize).into(),
            (ebss as usize).into(),
            MapType::Identical,
            MapPermission::R | MapPermission::W,
        ), None);

        info!("mapping physical memory");
        memory_set.push(MapArea::new(
            (ekernel as usize).into(),
            MEMORY_END.into(),
            MapType::Identical,
            MapPermission::R | MapPermission::W,
        ), None);

        memory_set
    }

    // 用于通过应用的 ELF 格式可执行文件 解析出各数据段并对应生成应用的地址空间
    /// Include sections in elf and trampoline and TrapContext and user stack,
    /// also returns user_sp and entry point.
    pub fn from_elf(elf_data: &[u8]) -> (Self, usize, usize) { 
        let mut memory_set = Self::new_bare();
        // 为跳板建立映射
        memory_set.map_trampoline();
        // 为程序的各个段建立映射 ( U flag )
        let elf = xmas_elf::ElfFile::new(elf_data).unwrap();
        let elf_header = elf.header;
        let magic = elf_header.pt1.magic;
        // 通过 elf 魔数判断是否为合法的 elf 文件
        assert_eq!(magic, [0x7f, 0x45, 0x4c, 0x46], "invalid elf!");
        // 获取 program header 的数量
        // Program header 描述的是一个段在文件中的位置、大小以及它被放进内存后所在的位置和大小等属性 详见博客(https://www.cnblogs.com/jiqingwu/p/elf_explore_2.html)。
        let ph_count = elf_header.pt2.ph_count();
        let mut max_end_vpn = VirtPageNum(0);
        for i in 0..ph_count {
            let ph = elf.program_header(i).unwrap();
            // program header 的类型是 LOAD 这表明它有被内核加载的必要
            // 此时不必理会其他类型的 program header 
            if ph.get_type().unwrap() == xmas_elf::program::Type::Load {
                // ph.virtual_addr() 获取起始地址
                // ph.mem_size() 获取大小
                let start_va: VirtAddr = (ph.virtual_addr() as usize).into();
                let end_va: VirtAddr = ((ph.virtual_addr() + ph.mem_size()) as usize).into();
                let mut map_perm = MapPermission::U;
                let ph_flags = ph.flags();
                if ph_flags.is_read() { map_perm |= MapPermission::R; }
                if ph_flags.is_write() { map_perm |= MapPermission::W; }
                if ph_flags.is_execute() { map_perm |= MapPermission::X; }
                let map_area = MapArea::new(
                    start_va,
                    end_va,
                    MapType::Framed,
                    map_perm,

                );
                // 维护 max_end_vpn 记录目前涉及到的最大的虚拟页号
                max_end_vpn = map_area.vpn_range.get_end();
                memory_set.push(
                    map_area,
                    Some(&elf.input[ph.offset() as usize..(ph.offset() + ph.file_size()) as usize])
                );
            }
        }
        // 为用户栈建立映射
        // 在 max_end_vpn `上面` 再放置一个保护页面和用户栈即可
        // 详见 file:///./app-as-full.png
        let max_end_va: VirtAddr = max_end_vpn.into();
        let mut user_stack_bottom: usize = max_end_va.into();

        // guard page 
        user_stack_bottom += PAGE_SIZE;
        let user_stack_top = user_stack_bottom + USER_STACK_SIZE;
        memory_set.push(MapArea::new(
            user_stack_bottom.into(),
            user_stack_top.into(),
            MapType::Framed,
            MapPermission::R | MapPermission::W | MapPermission::U,
        ), None);

        // 为 TrapContext 建立映射 (file:///./应用地址空间.md)
        memory_set.push(MapArea::new(
            TRAP_CONTEXT.into(),
            TRAMPOLINE.into(),
            MapType::Framed,
            MapPermission::R | MapPermission::W,
        ), None);
        // user_stack_top: 用户栈虚拟地址
        // elf.header.pt2.entry_point() 获取应用入口地址
        // 这两个参数将被用来创建应用的任务控制块
        (memory_set, user_stack_top, elf.header.pt2.entry_point() as usize)
    }

    pub fn activate(&self) {
        let satp = self.page_table.token();
        unsafe {
            // 在 satp寄存器 中写入多级页表的根结点
            satp::write(satp);
            // 从这一刻开始 SV39 分页模式就被启用了
            // 而且 MMU 会使用内核地址空间的多级页表进行地址转换
            // 我们必须注意切换 satp CSR 是否是一个 平滑 的过渡
            // 切换 satp 的指令及其下一条指令这两条相邻的指令的 虚拟地址是相邻的
            // （由于切换 satp 的指令并不是一条跳转指令， pc 只是简单的自增当前指令的字长）
            // 而它们所在的物理地址一般情况下也是相邻的
            // 但是它们所经过的地址转换流程却是不同的
            // 切换 satp 导致 MMU 查的多级页表 是不同的
            // 
            // 而我们采取的映射方式是恒等映射
            // 切换之前是视为物理地址直接取指，也可以将其看成一个恒等映射
            // 所以: 即使切换了地址空间，指令仍应该 能够被连续的执行

            // 刷新 快表 (TLB, Translation Lookaside Buffer)
            // 详见 file:///./SV39分页机制.md
            llvm_asm!("sfence.vma" :::: "volatile");
        }
    }
}


// 以逻辑段 MapArea 为单位描述一段连续地址的虚拟内存
// 所谓逻辑段，就是指地址区间中的一段实际可用（即 MMU 通过查多级页表 可以正确完成地址转换）的地址连续的虚拟地址区间
// 该区间内包含的所有虚拟页面都以一种相同的方式映射到物理页帧，具有可读/可写/可执行等属性
pub struct MapArea {
    vpn_range: VPNRange, // 描述一段虚拟页号的连续区间，表示该逻辑段在地址区间中的位置和长度
    data_frames: BTreeMap<VirtPageNum, FrameTracker>,
    map_type: MapType,
    map_perm: MapPermission
}

impl MapArea {
    // 用于新建一个逻辑段结构体
    pub fn new(
        start_va: VirtAddr,
        end_va: VirtAddr,
        map_type: MapType,
        map_perm: MapPermission
    ) -> Self {
        // 传入的起始/终止虚拟地址分别被下取整/上取整为虚拟页号并传入 迭代器 vpn_range 中
        let start_vpn: VirtPageNum = start_va.floor();
        let end_vpn: VirtPageNum = end_va.ceil();
        Self {
            vpn_range: VPNRange::new(start_vpn, end_vpn),
            data_frames: BTreeMap::new(),
            map_type,
            map_perm,
        }
    }

    // 为 vpn 在 page_table 中建立 vpn -> ppn(PTE) 的映射
    pub fn map_one(&mut self, page_table: &mut PageTable, vpn: VirtPageNum) {
        let ppn: PhysPageNum;

        // 根据该逻辑段的映射方式以及 vpn 决定 ppn
        match self.map_type {
            MapType::Identical => {
                ppn = PhysPageNum(vpn.0); // 恒等映射
            }
            MapType::Framed => {
                // 以 Framed 方式映射的时候，需要分配一个物理页帧作为当前的虚拟页面的映射
                let frame = frame_alloc().unwrap();
                ppn = frame.ppn;
                // 将这个物理页帧挂在逻辑段的 data_frames 字段下对其生命周期进行管理
                self.data_frames.insert(vpn, frame);
            }
        }
        // 根据 self.map_perm(MapPermission类型) 得到页表项的标志位 PTEFlags
        let pte_flags = PTEFlags::from_bits(self.map_perm.bits).unwrap();
        // ppn + pte_flags => PTE
        // 调用 PageTable 的 map 接口为 vpn -> ppn(PTE) 建立映射
        page_table.map(vpn, ppn, pte_flags);
    }

    // 将 vpn 在 page_table 中建立的 vpn -> ppn(PTE) 的映射移除
    pub fn unmap_one(&mut self, page_table: &mut PageTable, vpn: VirtPageNum) {
        match self.map_type {
            MapType::Framed => {
                // 这时候要注意对 vpn 对应的 物理帧 的生命周期进行管理
                self.data_frames.remove(&vpn);
            }
            _ => {}
        }
        // 调用 PageTable 的 unmap 接口就行
        page_table.unmap(vpn);
    }

    // 将当前逻辑段到物理内存的映射(体现在 self.vpn_range) 加入 传入的该逻辑段所属的地址空间的多级页表中
    pub fn map(&mut self, page_table: &mut PageTable) {
        for vpn in self.vpn_range {
            self.map_one(page_table, vpn);
        }
    }

    // 将当前逻辑段到物理内存的映射(体现在 self.vpn_range) 从 传入的该逻辑段所属的地址空间的多级页表中 删除
    pub fn unmap(&mut self, page_table: &mut PageTable) {
        for vpn in self.vpn_range {
            self.unmap_one(page_table, vpn);
        }
    }

    // 将切片 data 中的数据拷贝到当前逻辑段实际被内核放置在的各物理页帧上
    // 切片 data 中的数据大小不超过当前逻辑段的 总大小
    // 切片中的数据会被对齐到逻辑段的开头，然后逐页拷贝到实际的物理页帧
    pub fn copy_data(&mut self, page_table: &mut PageTable, data: &[u8]) {
        assert_eq!(self.map_type, MapType::Framed); // 通过断言确保映射方式为 Framed
        
        let mut start: usize = 0;
        let mut current_vpn = self.vpn_range.get_start();
        let len = data.len();
        loop {
            // 一个 PAGE_SIZE(或剩余部分) 大小的切片
            let src = &data[start..len.min(start + PAGE_SIZE)];
            let dst = &mut page_table
                .translate(current_vpn) // current_vpn -> PTE
                .unwrap()   // TODO: 考虑内存耗尽的情况
                .ppn()         // 获取 PhysPageNum
                .get_bytes_array()[..src.len()]; // 获取 PhysPageNum 对应的物理帧
            dst.copy_from_slice(src);   // 将 src 复制到该物理帧中
            start += PAGE_SIZE;
            if start >= len {
                break;
            }
            current_vpn.step(); // 若没复制完, 到下一个地方复制, 这样做的话, 数据也是连续的
        }
    }
}


// MapType 描述逻辑段内的所有虚拟页面映射到物理页帧的同一种方式
#[derive(Copy, Clone, PartialEq, Debug)]
pub enum MapType {
    Identical, // 恒等映射
    Framed,    // 对于每个虚拟页面都需要映射到一个新分配的物理页帧
    // 当逻辑段采用 MapType::Framed 方式映射到物理内存的时候
    // data_frames 是一个保存了该逻辑段内的每个虚拟页面 和它被映射到的物理页帧 FrameTracker 的一个键值对容器 BTreeMap 中，这些物理页帧被用来存放实际内存数据而不是 作为多级页表中的中间节点。
}


// 表示控制该逻辑段的访问方式
// 是页表项标志位 PTEFlags 的一个子集
// 仅保留 U/R/W/X 四个标志位
// 因为其他的标志位仅与硬件的地址转换机制细节相关
// 这样的设计能避免引入错误的标志位
bitflags! {
    pub struct MapPermission: u8 {
        const R = 1 << 1;
        const W = 1 << 2;
        const X = 1 << 3;
        const U = 1 << 4;
    }
}

// 我们既需要 Arc<T> 提供的共享 引用，也需要 Mutex<T> 提供的互斥访问
// 在多核环境下才能体现出它的全部能力，目前在单核环境下主要是为了通过编译器检查
lazy_static! {
    pub static ref KERNEL_SPACE: Arc<Mutex<MemorySet>> = Arc::new(Mutex::new(
        MemorySet::new_kernel()
    ));
}

pub fn remap_test() {
    let mut kernel_space = KERNEL_SPACE.lock();
    let mid_text: VirtAddr = ((stext as usize + etext as usize ) / 2).into();
    let mid_rodata: VirtAddr = ((srodata as usize + erodata as usize ) / 2).into();
    let mid_data: VirtAddr = ((sdata as usize + edata as usize ) / 2).into();
    assert_eq!(
        kernel_space.page_table.translate(mid_text.floor()).unwrap().writable(),
        false
    );
    assert_eq!(
        kernel_space.page_table.translate(mid_rodata.floor()).unwrap().writable(),
        false,
    );
    assert_eq!(
        kernel_space.page_table.translate(mid_data.floor()).unwrap().executable(),
        false,
    );
    info!("remap_test passed!");
}